from jax import random
from jax import vmap
from jax import jit
from jax import device_put
from jax import numpy as jnp
from PyCmpltrtok.common import sep, MyTimer


key = random.PRNGKey(0)
mat = random.normal(key, (150, 100))


def batched_x():
    x = random.normal(key, (1000, 100))
    return x


def apply_matrix(v):
    return jnp.dot(mat, v)


sep('apply_matrix with loop')
xN = 5


def naively_batched_apply_matrix(v_batched):
    res = jnp.stack([apply_matrix(v) for v in v_batched])
    return res


for _ in range(xN):
    bx = device_put(batched_x())
    with MyTimer() as timer:
        naively_batched_apply_matrix(bx)

sep('jit(apply_matrix with loop)')
jit_naively_batched_apply_matrix = jit(naively_batched_apply_matrix)
for _ in range(xN):
    bx = device_put(batched_x())
    with MyTimer() as timer:
        jit_naively_batched_apply_matrix(bx)

sep('vmap')
vmap_apply_matrix = vmap(apply_matrix)


def vmap_batched_apply_matrix(v_batched):
    return vmap_apply_matrix(v_batched)


for _ in range(xN):
    bx = device_put(batched_x())
    with MyTimer() as timer:
        vmap_batched_apply_matrix(bx)
